using System.Collections.Generic;
using System.Linq;
using System.Windows.Media;

namespace System.Windows
{
    /// <summary>
    /// <see cref="VisualTreeHelper">VisualTreeHelper</see> extensions.
    /// </summary>
    public static class VisualTreeHelperExtensions
    {
        /// <summary>
        /// Finds a Child of a given item in the visual tree. 
        /// </summary>
        /// <param name="parent">A direct parent of the queried item.</param>
        /// <typeparam name="T">The type of the queried item.</typeparam>
        /// <param name="childName">x:Name or Name of child. </param>
        /// <returns>The first parent item that matches the submitted type parameter. 
        /// If not matching item can be found, a null parent is being returned.</returns>
        public static T FindVisualChild<T>(this DependencyObject parent, string childName)
           where T : DependencyObject
        {
            if (parent == null)
            {
                return null;
            }

            T foundChild = null;

            var childrenCount = VisualTreeHelper.GetChildrenCount(parent);
            
            for (var i = 0; i < childrenCount; i++)
            {
                var child = VisualTreeHelper.GetChild(parent, i);
                var childType = child as T;
                
                if (childType == null)
                {
                    foundChild = FindVisualChild<T>(child, childName);

                    if (foundChild != null) break;
                }
                else if (!string.IsNullOrEmpty(childName))
                {
                    var frameworkElement = child as FrameworkElement;
                    
                    if (frameworkElement != null && frameworkElement.Name == childName)
                    {
                        foundChild = (T)child;
                        break;
                    }
                    
                    foundChild = FindVisualChild<T>(child, childName);
                    
                    if (foundChild != null) break;
                }
                else
                {
                    foundChild = (T)child;
                    break;
                }
            }

            return foundChild;
        }

        /// <summary>
        /// Finds a parent of a given item on the visual tree.
        /// </summary>
        /// <typeparam name="T">The type of the queried item.</typeparam>
        /// <param name="child">A direct or indirect child of the queried item.</param>
        /// <param name="name">Parent name</param>
        /// <returns>The first parent item that matches the submitted type parameter. 
        /// If not matching item can be found, a null reference is being returned.</returns>
        public static T FindVisualParent<T>(this DependencyObject child, string name = null)
          where T : DependencyObject
        {
            var parentObject = VisualTreeHelper.GetParent(child);

            if (parentObject == null)
            {
                return null;
            }

            var parent = parentObject as T;
            
            if (parent != null)
            {
                if (string.IsNullOrEmpty(name))
                {
                    return parent;
                }

                var childFrameworkElement = child as FrameworkElement;
                
                if (childFrameworkElement != null && name == childFrameworkElement.Name)
                {
                    return parent;
                }
            }
            
            return FindVisualParent<T>(parentObject, name);
        }

        /// <summary>
        /// Gets the visual parent of the element
        /// </summary>
        /// <param name="child">The element to check</param>
        /// <returns>The visual parent</returns>
        public static FrameworkElement FindVisualParent(this FrameworkElement child)
        {
            return VisualTreeHelper.GetParent(child) as FrameworkElement;
        }

        /// <summary>
        /// Gets the ancestors of the element, up to the root
        /// </summary>
        /// <param name="element">The element to start from</param>
        /// <returns>An enumerator of the ancestors</returns>
        public static IEnumerable<FrameworkElement> GetVisualAncestors(this FrameworkElement element)
        {
            var parent = element.FindVisualParent();

            while (parent != null)
            {
                yield return parent;

                parent = parent.FindVisualParent();
            }
        }

        /// <summary>
        /// Gets the first ancestor of the element. If nothing found returns null.
        /// </summary>
        /// <param name="element">The element to start from</param>
        /// <returns>Ancestor or null.</returns>
        public static T GetFirstOrDefaultVisualAncestor<T>(this FrameworkElement element) where T : class
        {
            if (element == null)
            {
                throw new ArgumentNullException("element");
            }

            return element.GetVisualAncestors().FirstOrDefault(c => c is T) as T;
        }

        /// <summary>
        /// Gets all visuals as <see cref="IEnumerable{T}">enumeration of</see> <see cref="DependencyObject">dependency objects</see> for provided root <see cref="DependencyObject">dependency object</see>.
        /// </summary>
        /// <param name="root">Root <see cref="DependencyObject">dependency object</see>.</param>
        /// <returns><see cref="IEnumerable{T}">Enumeration of</see> <see cref="DependencyObject">dependency objects</see>.</returns>
        public static IEnumerable<DependencyObject> GetVisuals(this DependencyObject root)
        {
            var count = VisualTreeHelper.GetChildrenCount(root);
            
            for (var i = 0; i < count; i++)
            {
                var child = VisualTreeHelper.GetChild(root, i);
                yield return child;

                foreach (var descendants in child.GetVisuals())
                {
                    yield return descendants;
                }
            }
        }

        /// <summary>
        /// Gets <see cref="IEnumerable{T}">enumeration of</see> child objects for provided root <see cref="DependencyObject">dependency object</see>.
        /// </summary>
        /// <param name="obj">Provided root <see cref="DependencyObject">dependency object</see>.</param>
        /// <param name="name">Optional: name of the child object.</param>
        /// <typeparam name="T">Type of the child objects to be found.</typeparam>
        /// <returns><see cref="IEnumerable{T}">Enumeration</see> of child objects.</returns>
        public static IEnumerable<T> GetChildObjects<T>(this DependencyObject obj, string name = null)
        {
            var retVal = new List<T>();

            for (var i = 0; i < VisualTreeHelper.GetChildrenCount(obj); i++)
            {
                object c = VisualTreeHelper.GetChild(obj, i);

                if (c.GetType().FullName == typeof (T).FullName
                    && (string.IsNullOrEmpty(name) || ((FrameworkElement) c).Name == name))
                {
                    retVal.Add((T) c);
                }

                var gc = ((DependencyObject) c).GetChildObjects<T>(name);

                if (gc != null)
                {
                    retVal.AddRange(gc);
                }
            }

            return retVal;
        }
    }
}